#! /usr/bin/env python
# coding=utf-8



from Bio import SeqIO
import os
import shutil

import sys
import argparse
from argparse import RawTextHelpFormatter

parser = argparse.ArgumentParser(
    description='''
    对NGSpeciesID的结果中 找出指定标本 含有指定reads数以上的cluster
    用法:
    barcode_filter_by_reads_and_specimen -i ngsp_out -l specimen_lista -r 25 -o ngsp_filter_out
    由大天才于2022年3月24日创建于浙江农业大学''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-i',
                help='必须给定，run_NGSpeciesID.py的输出路径')

parser.add_argument('-l',
                help='必须给定, 标本列表')

parser.add_argument('-r',
                help='必须给定, reads数')

parser.add_argument('-o',
                help='必须给定，输出路径')





# 通过reads 数 和 标本 id号 过滤结果


# lista_file = r'specimen_lista'

# reads_cutoff = 25

# new_path = r'ngsp_filter_out'

# ngsp_out_path = r'ngsp_out'



args = parser.parse_args()

if not args.i or not args.o or not args.r or not args.l:

    parser.print_help()

    sys.exit()

reads_cutoff = int(args.r)



new_path = args.o #'ngsp_filter_out'

ngsp_out_path = args.i

lista_file  =args.l




try:
	shutil.rmtree(new_path)
except:
	pass



try:
	os.mkdir(new_path)
except:
	pass
	

# 读取标本列表
spe_lista = []

with open(lista_file) as fila:
	for i in fila:
		k = i.strip()
		if k!='':
			spe_lista.append(k)

reads_nub_pass_cluster =0
reads_nub_pass_specimen = 0

for i in os.listdir(ngsp_out_path):
	p = ngsp_out_path + '/' + i
	if i in spe_lista:
		for j in os.listdir(p):
			if j.split('.')[-1] == 'fasta':
				fasta_p = p +'/' + j
				s = next(SeqIO.parse(fasta_p,'fasta'))
				reads_nub = int(s.name.split('_')[-1])
				if reads_nub >= reads_cutoff:

					if not os.path.exists(new_path+'/'+i):
						os.mkdir(new_path+'/'+i)
						reads_nub_pass_specimen += 1
					with open(new_path+'/'+i+'/'+j,'w') as fita:
						fita.write('>'+s.name+'\n'+str(s.seq)+'\n') 
					reads_nub_pass_cluster +=1 



print('一共过滤得到 %s 个cluster 来自 %s 个标本' %(reads_nub_pass_cluster, reads_nub_pass_specimen ))


				